from tensorflow.python.platform import gfile
import tensorflow as tf
import os
import cv2
import numpy as np
from utiles.box_utiles import polygons_from_bitmap
from utiles.transform import resize_image

IMAGE_SIZE = [160, 160]
pb_file_path = './pbMode/123.pb'
image_root = 'E:/CODES/DifferentiableBinarization/datasets/selfmade3/val_images/'

sess = tf.Session()
with gfile.FastGFile(pb_file_path, 'rb') as f:
	graph_def = tf.GraphDef()
	graph_def.ParseFromString(f.read())
	sess.graph.as_default()
	tf.import_graph_def(graph_def, name='')
sess.run(tf.global_variables_initializer())

decision_out = sess.graph.get_tensor_by_name('dbnet/proba3_sigmoid:0')
input_image = sess.graph.get_tensor_by_name('image_input:0')
image_names = [i[2] for i in os.walk(image_root)][0]
name_dic = {}
count = 0
for i in image_names:
	image_path = os.path.join(image_root, i)
	image = cv2.imread(image_path)
	image = image.astype(np.float32)

	image = resize_image(IMAGE_SIZE[0], image)
	image = np.array(image[np.newaxis,:, :, :])
	image /= 255.0
	# mean = [103.939, 116.779, 123.68]
	# image[..., 0] -= mean[0]
	# image[..., 1] -= mean[1]
	# image[..., 2] -= mean[2]
	proba_map = sess.run(decision_out, feed_dict={input_image: image})

	proba_map = np.squeeze(proba_map, 0)
	image = np.squeeze(image, 0)
	image *= 255.0

	# image[..., 0] += mean[0]
	# image[..., 1] += mean[1]
	# image[..., 2] += mean[2]
	img = image.copy()
	bitmap = proba_map > 0.3
	boxes, scores = polygons_from_bitmap(proba_map, bitmap, image.shape[0], image.shape[1], box_thresh=0.5)
	# for box in boxes:
	# 	cv2.drawContours(image, [np.array(box)], -1, (0, 255, 0), 2)

	# cv2.imwrite(os.path.join('./test/', i), image)

	for contour in boxes:
		contour = np.reshape(contour, (-1, 2))
		xmin = sorted(list(contour), key=lambda x: x[0])[0][0]
		xmax = sorted(list(contour), key=lambda x: x[0])[-1][0]
		ymin = sorted(list(contour), key=lambda x: x[1])[0][1]
		ymax = sorted(list(contour), key=lambda x: x[1])[-1][1]
		cut_image = img[ymin:ymax, xmin:xmax, :]

		cv2.imwrite(os.path.join('./cut/', str(count) + '-' + str(count) + '.jpg'), cut_image)
		count += 1






